{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Streaming"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain import hub\n",
    "from langchain.agents import AgentExecutor, create_openai_tools_agent\n",
    "from langchain.tools import tool\n",
    "from langchain_core.callbacks import Callbacks\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_openai import ChatOpenAI\n",
    "import random\n",
    "\n",
    "\n",
    "model = ChatOpenAI(temperature=0, streaming=True)\n",
    "\n",
    "@tool\n",
    "async def where_cat_is_hiding() -> str:\n",
    "    \"\"\"Where is the cat hiding right now?\"\"\"\n",
    "    return random.choice([\"under the bed\", \"on the shelf\"])\n",
    "\n",
    "\n",
    "@tool\n",
    "async def get_items(place: str) -> str:\n",
    "    \"\"\"Use this tool to look up which items are in the given place.\"\"\"\n",
    "    if \"bed\" in place:  # For under the bed\n",
    "        return \"socks, shoes and dust bunnies\"\n",
    "    if \"shelf\" in place:  # For 'shelf'\n",
    "        return \"books, penciles and pictures\"\n",
    "    else:  # if the agent decides to ask about a different place\n",
    "        return \"cat snacks\"\n",
    "\n",
    "# Get the prompt to use - you can modify this!\n",
    "prompt = hub.pull(\"hwchase17/openai-tools-agent\")\n",
    "# print(prompt.messages) -- to see the prompt\n",
    "tools = [get_items, where_cat_is_hiding]\n",
    "\n",
    "\n",
    "agent = create_openai_tools_agent(\n",
    "    model.with_config({\"tags\": [\"agent_llm\"]}), tools, prompt\n",
    ")\n",
    "#  associated the name Agent with our agent using \"run_name\"=\"Agent\".\n",
    "agent_executor = AgentExecutor(agent=agent, tools=tools).with_config(\n",
    "    {\"run_name\": \"Agent\"}\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "------\n",
      "{'actions': [...], 'messages': [...]}\n",
      "------\n",
      "{'messages': [...], 'steps': [...]}\n",
      "------\n",
      "{'actions': [...], 'messages': [...]}\n",
      "------\n",
      "{'messages': [...], 'steps': [...]}\n",
      "------\n",
      "{'messages': [...],\n",
      " 'output': 'The items located where the cat is hiding (under the bed) are '\n",
      "           'socks, shoes, and dust bunnies.'}\n"
     ]
    }
   ],
   "source": [
    "# Note: We use `pprint` to print only to depth 1, it makes it easier to see the output from a high level, before digging in.\n",
    "import pprint\n",
    "\n",
    "chunks = []\n",
    "\n",
    "async for chunk in agent_executor.astream(\n",
    "    {\"input\": \"what's items are located where the cat is hiding?\"}\n",
    "):\n",
    "    chunks.append(chunk)\n",
    "    print(\"------\")\n",
    "    pprint.pprint(chunk, depth=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ToolAgentAction(tool='where_cat_is_hiding', tool_input={}, log='\\nInvoking: `where_cat_is_hiding` with `{}`\\n\\n\\n', message_log=[AIMessageChunk(content='', additional_kwargs={'tool_calls': [{'index': 0, 'id': 'call_19WpkitCG60CQgiNeHKuXzFH', 'function': {'arguments': '{}', 'name': 'where_cat_is_hiding'}, 'type': 'function'}]}, response_metadata={'finish_reason': 'tool_calls'}, id='run-fd281b15-5a3f-41e3-8fb6-de6e60bcbc5e', tool_calls=[{'name': 'where_cat_is_hiding', 'args': {}, 'id': 'call_19WpkitCG60CQgiNeHKuXzFH'}], tool_call_chunks=[{'name': 'where_cat_is_hiding', 'args': '{}', 'id': 'call_19WpkitCG60CQgiNeHKuXzFH', 'index': 0}])], tool_call_id='call_19WpkitCG60CQgiNeHKuXzFH')]\n",
      "[AIMessageChunk(content='', additional_kwargs={'tool_calls': [{'index': 0, 'id': 'call_19WpkitCG60CQgiNeHKuXzFH', 'function': {'arguments': '{}', 'name': 'where_cat_is_hiding'}, 'type': 'function'}]}, response_metadata={'finish_reason': 'tool_calls'}, id='run-fd281b15-5a3f-41e3-8fb6-de6e60bcbc5e', tool_calls=[{'name': 'where_cat_is_hiding', 'args': {}, 'id': 'call_19WpkitCG60CQgiNeHKuXzFH'}], tool_call_chunks=[{'name': 'where_cat_is_hiding', 'args': '{}', 'id': 'call_19WpkitCG60CQgiNeHKuXzFH', 'index': 0}])]\n",
      "[FunctionMessage(content='under the bed', name='where_cat_is_hiding')]\n",
      "[AIMessageChunk(content='', additional_kwargs={'tool_calls': [{'index': 0, 'id': 'call_iZnTjNpl9tbbAADn5IawhcgW', 'function': {'arguments': '{\"place\":\"under the bed\"}', 'name': 'get_items'}, 'type': 'function'}]}, response_metadata={'finish_reason': 'tool_calls'}, id='run-36fe46c8-82ac-42dc-9a18-489cfcbea2ad', tool_calls=[{'name': 'get_items', 'args': {'place': 'under the bed'}, 'id': 'call_iZnTjNpl9tbbAADn5IawhcgW'}], tool_call_chunks=[{'name': 'get_items', 'args': '{\"place\":\"under the bed\"}', 'id': 'call_iZnTjNpl9tbbAADn5IawhcgW', 'index': 0}])]\n",
      "[FunctionMessage(content='socks, shoes and dust bunnies', name='get_items')]\n",
      "[AIMessage(content='The items located where the cat is hiding (under the bed) are socks, shoes, and dust bunnies.')]\n"
     ]
    }
   ],
   "source": [
    "print(chunks[0][\"actions\"])\n",
    "for chunk in chunks:\n",
    "    print(chunk[\"messages\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "async for chunk in agent_executor.astream(\n",
    "    {\"input\": \"what's items are located where the cat is hiding?\"}\n",
    "):\n",
    "    # Agent Action\n",
    "    if \"actions\" in chunk:\n",
    "        for action in chunk[\"actions\"]:\n",
    "            print(f\"Calling Tool: `{action.tool}` with input `{action.tool_input}`\")\n",
    "    # Observation\n",
    "    elif \"steps\" in chunk:\n",
    "        for step in chunk[\"steps\"]:\n",
    "            print(f\"Tool Result: `{step.observation}`\")\n",
    "    # Final result\n",
    "    elif \"output\" in chunk:\n",
    "        print(f'Final Output: {chunk[\"output\"]}')\n",
    "    else:\n",
    "        raise ValueError()\n",
    "    print(\"---\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain0_1",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
